import os
from os.path import join
import cv2
import numpy as np
import open3d as o3d


def txt2intrinsic(intrinsic_path):
    I = []
    dist = None
    with open(intrinsic_path, "r") as f:
        for line in f:
            values = [float(v) for v in line.strip().split(",")]
            if len(values) == 3:
                I.append(values)
            else:
                dist = values
    I = np.float32(I)
    dist = np.float32(dist)
    return I, dist


def visualize_world_rgbd(rgb_video_path, depth_dir, intrinsic_mat, cam_pose, frame_idx=0, mask_path=None):
    cap = cv2.VideoCapture(rgb_video_path)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    cap.set(cv2.CAP_PROP_FOURCC, fourcc)

    rgb = None
    frame_cnt = -1
    while True:
        frame_cnt += 1
        suc, img = cap.read()
        if not suc:
            break
        if frame_cnt == int(frame_idx):
            img = img[:, :, ::-1].astype(np.uint8)
            rgb = o3d.geometry.Image(img)
            break
    assert not rgb is None

    depth = o3d.io.read_image(join(depth_dir, str(frame_idx) + ".png"))

    if not mask_path is None:
        cap = cv2.VideoCapture(mask_path)
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        cap.set(cv2.CAP_PROP_FOURCC, fourcc)
        suc = cap.isOpened()
        frame_cnt = -1
        while True:
            frame_cnt += 1
            suc, img = cap.read()
            if not suc:
                break
            if frame_cnt == int(frame_idx):
                mask = img[:, :, ::-1].astype(np.uint8)
                break
        mask = mask[..., 0] > 10  # fake
        depth_np = np.array(depth)
        depth_np[~mask] = 0
        depth = o3d.geometry.Image(depth_np)

    intrinsic = o3d.camera.PinholeCameraIntrinsic(1280, 720, intrinsic_mat[0, 0], intrinsic_mat[1, 1], intrinsic_mat[0, 2], intrinsic_mat[1, 2])
    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(rgb, depth, convert_rgb_to_intensity=False, depth_trunc=10.0)
    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, intrinsic)

    # cam2world
    points = np.float32(pcd.points)
    points = points @ cam_pose[:3, :3].T + cam_pose[:3, 3]
    pcd.points = o3d.utility.Vector3dVector(points)

    return pcd


if __name__ == "__main__":
    data_dir = "./data/20230630_test_3views"
    save_dir = "./data/20230630_test_mv_pcd"

    os.makedirs(save_dir, exist_ok=True)

    label = "_d455_camera1"
    intrinsic_path = "../collect_rawdata/camera_calib/d455_1/intrinsic.txt"
    cam_pose_path = "../collect_rawdata/camera_calib/d455_1/camera2world.txt"
    intrinsic_mat, _ = txt2intrinsic(intrinsic_path)
    cam_pose = np.loadtxt(cam_pose_path)
    os.makedirs(join(data_dir, "decode" + label), exist_ok=True)
    os.system("ffmpeg -i {} -f image2 -start_number 0 -vf fps=fps=15 -qscale:v 2 ./{}/%d.png -loglevel quiet".format(join(data_dir, label + "_aligned_depth_to_color_image_raw.avi"), join(data_dir, "decode" + label)))
    pcd = visualize_world_rgbd(join(data_dir, label + "_color_image_raw.mp4"), join(data_dir, "decode" + label), intrinsic_mat, cam_pose, frame_idx=0, mask_path=None)
    o3d.io.write_point_cloud(join(save_dir, label + ".ply"), pcd)

    label = "_d455_camera2"
    intrinsic_path = "../collect_rawdata/camera_calib/d455_2/intrinsic.txt"
    cam_pose_path = "../collect_rawdata/camera_calib/d455_2/camera2world.txt"
    intrinsic_mat, _ = txt2intrinsic(intrinsic_path)
    cam_pose = np.loadtxt(cam_pose_path)
    os.makedirs(join(data_dir, "decode" + label), exist_ok=True)
    os.system("ffmpeg -i {} -f image2 -start_number 0 -vf fps=fps=15 -qscale:v 2 ./{}/%d.png -loglevel quiet".format(join(data_dir, label + "_aligned_depth_to_color_image_raw.avi"), join(data_dir, "decode" + label)))
    pcd = visualize_world_rgbd(join(data_dir, label + "_color_image_raw.mp4"), join(data_dir, "decode" + label), intrinsic_mat, cam_pose, frame_idx=0, mask_path=None)
    o3d.io.write_point_cloud(join(save_dir, label + ".ply"), pcd)
